from .parser import AST
from .juliet_adaption import juliet_file_filter
import os, random


def data_augmentation_main(vulnerable_code_dataset: str, secure_code_dataset: str, augmented_code_dir: str) -> None:
    def get_all_filenames(directory) -> list:
        return [os.path.join(root, file) for root, _, files in os.walk(directory) for file in files]

    os.makedirs(augmented_code_dir, exist_ok=True)
    if not os.path.exists(vulnerable_code_dataset):
        raise FileNotFoundError(f"The directory {vulnerable_code_dataset} does not exist.")
    if not os.path.exists(secure_code_dataset):
        raise FileNotFoundError(f"The directory {secure_code_dataset} does not exist.")
    
    for root, _, _ in os.walk(vulnerable_code_dataset):
        relative_path = os.path.join(os.path.basename(vulnerable_code_dataset) + "_augmented", os.path.relpath(root, vulnerable_code_dataset))
        os.makedirs(os.path.join(augmented_code_dir, relative_path), exist_ok=True)
    
    for file in get_all_filenames(vulnerable_code_dataset):
        destination = os.path.join(augmented_code_dir, os.path.basename(vulnerable_code_dataset) + "_augmented", os.path.relpath(file, vulnerable_code_dataset))
        if juliet_file_filter(file):
            if augment(file, random.choice(get_all_filenames(secure_code_dataset)), destination):
                print(f"File {file} has been augmented.")
            else:
                print(f"Failed to augment file {file}.")
        else:
            with open(file, 'rb') as src_file, open(destination, 'wb') as dst_file:
                dst_file.write(src_file.read())
            print(f"File {file} dose not match condition, skipping...")
    return
    

def augment(vulnerable_code_file, secure_code_file, augmented_code_file) -> bool:
    with open(vulnerable_code_file,'rb') as v_f:
        vulnerable_code = v_f.read()
    with open(secure_code_file,'rb') as s_f:
        secure_code = s_f.read()
    vulnerable_ast = AST(vulnerable_code)
    secure_ast = AST(secure_code)
    vulnerable_ast.insert_point_get()
    secure_ast.code_split()
    if not vulnerable_ast.code_merge(secure_ast):
        return False
    vulnerable_ast.code_output(augmented_code_file)
    return True